#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "mbedtls/des.h"
#include "mbedtls/aes.h"

#include "abcrypto.h"
#include "cipher.h"

struct mbed_cipher_ctx {
	mbedtls_aes_context aes_ctx;
	uint16_t alg;
};

static int init_ctx(void **ctx, uint8_t *key, uint32_t klen, uint8_t *iv,
		    uint16_t alg, bool isenc)
{
	struct mbed_cipher_ctx *mbed_cipher_ctx = NULL;
	mbedtls_aes_context *aes_ctx;

	mbed_cipher_ctx = malloc(sizeof(struct mbed_cipher_ctx));
	if (mbed_cipher_ctx == NULL) {
		fprintf(stderr, "Alloc aesctx failed\n");
		return -1;
	}

	mbed_cipher_ctx->alg = alg;
	aes_ctx = &mbed_cipher_ctx->aes_ctx;

	mbedtls_aes_init(aes_ctx);

	switch (alg) {
	case ABCRYPTO_CIPHER_AES_128_ECB:
	case ABCRYPTO_CIPHER_AES_128_CBC:
	case ABCRYPTO_CIPHER_AES_128_CTR:
	case ABCRYPTO_CIPHER_AES_128_CFB:
	case ABCRYPTO_CIPHER_AES_128_OFB:
		if (klen != 16) {
			fprintf(stderr, "Invalid key len\n");
			mbedtls_aes_free(aes_ctx);
			return -1;
		}
		if (isenc)
			mbedtls_aes_setkey_enc(aes_ctx, key, 128);
		else
			mbedtls_aes_setkey_dec(aes_ctx, key, 128);

		break;
	case ABCRYPTO_CIPHER_AES_192_ECB:
	case ABCRYPTO_CIPHER_AES_192_CBC:
	case ABCRYPTO_CIPHER_AES_192_CTR:
	case ABCRYPTO_CIPHER_AES_192_CFB:
	case ABCRYPTO_CIPHER_AES_192_OFB:
		if (klen != 24) {
			fprintf(stderr, "Invalid key len\n");
			mbedtls_aes_free(aes_ctx);
			return -1;
		}
		if (isenc)
			mbedtls_aes_setkey_enc(aes_ctx, key, 192);
		else
			mbedtls_aes_setkey_dec(aes_ctx, key, 192);
		break;

	case ABCRYPTO_CIPHER_AES_256_ECB:
	case ABCRYPTO_CIPHER_AES_256_CBC:
	case ABCRYPTO_CIPHER_AES_256_CTR:
	case ABCRYPTO_CIPHER_AES_256_CFB:
	case ABCRYPTO_CIPHER_AES_256_OFB:
		if (klen != 32) {
			fprintf(stderr, "Invalid key len\n");
			mbedtls_aes_free(aes_ctx);
			return -1;
		}
		if (isenc)
			mbedtls_aes_setkey_enc(aes_ctx, key, 256);
		else
			mbedtls_aes_setkey_dec(aes_ctx, key, 256);
		break;
	default:
		fprintf(stderr, "Invalid alg\n");
		mbedtls_aes_free(aes_ctx);
		return -1;
	}

	*ctx = aes_ctx;

	return 0;
}

static void free_ctx(void *ctx)
{
	struct mbed_cipher_ctx *mbed_cipher_ctx = ctx;
	mbedtls_aes_free(&mbed_cipher_ctx->aes_ctx);
	free(ctx);
	return;
}

static int encrypt(void *ctx, uint8_t *iv, uint8_t *data, uint32_t len,
		   uint8_t *dst)
{
	struct mbed_cipher_ctx *mbed_cipher_ctx = ctx;

	int ret = -1;
	int i;
	size_t nc_off = 0;

	switch (mbed_cipher_ctx->alg & 0xf) {
	case ABCRYPTO_CIPHER_MODE_ECB:
		for (i = 0; i < len; i += 16) {
			ret = mbedtls_aes_crypt_ecb(&mbed_cipher_ctx->aes_ctx,
						    MBEDTLS_AES_ENCRYPT,
						    data + i, dst + i);
			if (ret) {
				fprintf(stderr, "ECB encrypt error\n");
				return -1;
			}
		}
		break;
	case ABCRYPTO_CIPHER_MODE_CBC:
		ret = mbedtls_aes_crypt_cbc(&mbed_cipher_ctx->aes_ctx,
					    MBEDTLS_AES_ENCRYPT, len, iv, data,
					    dst);

		break;
	case ABCRYPTO_CIPHER_MODE_CTR: {
		unsigned char stream_block[16];

		ret = mbedtls_aes_crypt_ctr(&mbed_cipher_ctx->aes_ctx, len,
					    &nc_off, iv, stream_block, data,
					    dst);
	}

	break;
	case ABCRYPTO_CIPHER_MODE_CFB:
		ret = mbedtls_aes_crypt_cfb128(&mbed_cipher_ctx->aes_ctx,
					       MBEDTLS_AES_ENCRYPT, len,
					       &nc_off, iv, data, dst);

		break;
	case ABCRYPTO_CIPHER_MODE_OFB:
		ret = mbedtls_aes_crypt_ofb(&mbed_cipher_ctx->aes_ctx, len,
					    &nc_off, iv, data, dst);
		break;
	default:
		fprintf(stderr, "Invalid cipher mode %u\n",
			mbed_cipher_ctx->alg & 0xf);
		break;
	}

	return ret;
}

static int decrypt(void *ctx, uint8_t *iv, uint8_t *data, uint32_t len,
		   uint8_t *dst)
{
	struct mbed_cipher_ctx *mbed_cipher_ctx = ctx;

	int ret = -1;
	int i;
	size_t nc_off = 0;

	switch (mbed_cipher_ctx->alg & 0xf) {
	case ABCRYPTO_CIPHER_MODE_ECB:
		for (i = 0; i < len; i += 16) {
			ret = mbedtls_aes_crypt_ecb(&mbed_cipher_ctx->aes_ctx,
						    MBEDTLS_AES_DECRYPT,
						    data + i, dst + i);
			if (ret) {
				fprintf(stderr, "ECB decrypt error\n");
				return -1;
			}
		}
		break;
	case ABCRYPTO_CIPHER_MODE_CBC:
		ret = mbedtls_aes_crypt_cbc(&mbed_cipher_ctx->aes_ctx,
					    MBEDTLS_AES_DECRYPT, len, iv, data,
					    dst);

		break;
	case ABCRYPTO_CIPHER_MODE_CTR: {
		unsigned char stream_block[16];

		ret = mbedtls_aes_crypt_ctr(&mbed_cipher_ctx->aes_ctx, len,
					    &nc_off, iv, stream_block, data,
					    dst);
	}

	break;
	case ABCRYPTO_CIPHER_MODE_CFB:
		ret = mbedtls_aes_crypt_cfb128(&mbed_cipher_ctx->aes_ctx,
					       MBEDTLS_AES_DECRYPT, len,
					       &nc_off, iv, data, dst);

		break;
	case ABCRYPTO_CIPHER_MODE_OFB:
		ret = mbedtls_aes_crypt_ofb(&mbed_cipher_ctx->aes_ctx, len,
					    &nc_off, iv, data, dst);
		break;
	default:
		fprintf(stderr, "Invalid cipher mode %u\n",
			mbed_cipher_ctx->alg & 0xf);
		break;
	}

	return ret;
}

static int once_encrypt(uint8_t *key, uint8_t *iv, uint8_t *src, uint8_t *dst,
			uint32_t len, uint16_t alg)
{
	void *ctx = NULL;
	int ret;
	uint32_t klen = 16;

	switch (alg) {
	case ABCRYPTO_CIPHER_AES_128_ECB:
	case ABCRYPTO_CIPHER_AES_128_CBC:
	case ABCRYPTO_CIPHER_AES_128_CTR:
	case ABCRYPTO_CIPHER_AES_128_CFB:
	case ABCRYPTO_CIPHER_AES_128_OFB:
		klen = 16;
		break;
	case ABCRYPTO_CIPHER_AES_192_ECB:
	case ABCRYPTO_CIPHER_AES_192_CBC:
	case ABCRYPTO_CIPHER_AES_192_CTR:
	case ABCRYPTO_CIPHER_AES_192_CFB:
	case ABCRYPTO_CIPHER_AES_192_OFB:
		klen = 24;
		break;
	case ABCRYPTO_CIPHER_AES_256_ECB:
	case ABCRYPTO_CIPHER_AES_256_CBC:
	case ABCRYPTO_CIPHER_AES_256_CTR:
	case ABCRYPTO_CIPHER_AES_256_CFB:
	case ABCRYPTO_CIPHER_AES_256_OFB:
		klen = 32;
		break;
	default:
		break;
	}
	ret = init_ctx(&ctx, key, klen, iv, alg, true);
	if (ret) {
		return -1;
	}

	ret = encrypt(ctx, iv, src, len, dst);

	free_ctx(ctx);
	return ret;
}

static int once_decrypt(uint8_t *key, uint8_t *iv, uint8_t *src, uint8_t *dst,
			uint32_t len, uint16_t alg)
{
	void *ctx = NULL;
	int ret;
	uint32_t klen = 16;

	switch (alg) {
	case ABCRYPTO_CIPHER_AES_128_ECB:
	case ABCRYPTO_CIPHER_AES_128_CBC:
	case ABCRYPTO_CIPHER_AES_128_CTR:
	case ABCRYPTO_CIPHER_AES_128_CFB:
	case ABCRYPTO_CIPHER_AES_128_OFB:
		klen = 16;
		break;
	case ABCRYPTO_CIPHER_AES_192_ECB:
	case ABCRYPTO_CIPHER_AES_192_CBC:
	case ABCRYPTO_CIPHER_AES_192_CTR:
	case ABCRYPTO_CIPHER_AES_192_CFB:
	case ABCRYPTO_CIPHER_AES_192_OFB:
		klen = 24;
		break;
	case ABCRYPTO_CIPHER_AES_256_ECB:
	case ABCRYPTO_CIPHER_AES_256_CBC:
	case ABCRYPTO_CIPHER_AES_256_CTR:
	case ABCRYPTO_CIPHER_AES_256_CFB:
	case ABCRYPTO_CIPHER_AES_256_OFB:
		klen = 32;
		break;
	default:
		break;
	}

	ret = init_ctx(&ctx, key, klen, iv, alg, false);
	if (ret) {
		return -1;
	}

	ret = decrypt(ctx, iv, src, len, dst);
	free_ctx(ctx);

	return ret;
}

static struct abcrypto_skcipher _mbed_cipher = { .init_ctx = init_ctx,
						 .free_ctx = free_ctx,
						 .encrypt = encrypt,
						 .decrypt = decrypt,
						 .once_encrypt = once_encrypt,
						 .once_decrypt = once_decrypt };

struct abcrypto_skcipher *mbed_cipher(void)
{
	return &_mbed_cipher;
}
